from django.core.management.base import BaseCommand
from django.contrib.auth import get_user_model
from django.contrib.auth.models import Group

User = get_user_model()

class Command(BaseCommand):
    help = '初始化测试数据'

    def handle(self, *args, **kwargs):
        self.stdout.write(self.style.SUCCESS('开始初始化测试数据...'))

        # 创建分析师组
        analyst_group, created = Group.objects.get_or_create(name='analysts')
        if created:
            self.stdout.write(self.style.SUCCESS('创建分析师组成功'))

        # 创建分析师用户
        analyst, created = User.objects.get_or_create(
            username='analyst1',
            defaults={
                'email': 'analyst1@example.com',
                'role': 'analyst',
                'is_staff': True
            }
        )
        if created:
            analyst.set_password('analyst123')
            analyst.save()
            analyst.groups.add(analyst_group)
            self.stdout.write(self.style.SUCCESS('创建分析师用户成功'))

        # 创建普通用户
        user, created = User.objects.get_or_create(
            username='user1',
            defaults={
                'email': 'user1@example.com',
                'role': 'user'
            }
        )
        if created:
            user.set_password('user123')
            user.save()
            self.stdout.write(self.style.SUCCESS('创建普通用户成功'))

        self.stdout.write(self.style.SUCCESS('测试数据初始化完成')) 